package my.algo.graph;

import java.util.LinkedList;
import java.util.PriorityQueue;

/**
 * 带权图的最小生成树
 */
public class LazyPrimMst {

    /** 标记顶点是否已加入 */
    private final boolean[] marked;
    private final LinkedList<Edge> queue;
    private final PriorityQueue<Edge> pq;

    public LazyPrimMst(EdgeWeightedGraph graph) {
        pq = new PriorityQueue<>();
        marked = new boolean[graph.V()];
        queue = new LinkedList<>();
        visit(graph, 0);
        while (!pq.isEmpty()) {
            Edge edge = pq.poll();
            int v = edge.either(), w = edge.other(v);
            if (marked[v] && marked[w]) {
                continue;
            }
            queue.add(edge);
            if (!marked[v]) {
                visit(graph, v);
            }
            if (!marked[w]) {
                visit(graph, w);
            }
        }
    }

    private void visit(EdgeWeightedGraph graph, int v) {
        marked[v] = true;
        for (Edge edge : graph.adj(v)) {
            if (!marked[edge.other(v)]) {
                pq.add(edge);
            }
        }
    }

    public Iterable<Edge> edges() {
        return queue;
    }


    public static void main(String[] args) {
        EdgeWeightedGraph graph = new EdgeWeightedGraph(9);
        graph.addEdge(new Edge(0, 1, 4));
        graph.addEdge(new Edge(0, 7, 8));
        graph.addEdge(new Edge(1, 2, 8));
        graph.addEdge(new Edge(1, 7, 11));
        graph.addEdge(new Edge(2, 3, 7));
        graph.addEdge(new Edge(2, 5, 4));
        graph.addEdge(new Edge(2, 8, 2));
        graph.addEdge(new Edge(3, 4, 9));
        graph.addEdge(new Edge(3, 5, 14));
        graph.addEdge(new Edge(4, 5, 10));
        graph.addEdge(new Edge(5, 6, 2));
        graph.addEdge(new Edge(6, 7, 1));
        graph.addEdge(new Edge(6, 8, 6));
        graph.addEdge(new Edge(7, 8, 7));
        LazyPrimMst lazyPrimMst = new LazyPrimMst(graph);
        Iterable<Edge> edges = lazyPrimMst.edges();
        System.out.println();
    }

}
